/**
 * \file encryption.cc
 * \namespace encryption
 *
 * Interfaces to encrypt and decrypt data with AES256 CBC having 
 * the password stored in a file. Uses OpenSSL routines.
 *
 * Developed by Jakob Blomer 2010 at CERN
 * jakob.blomer@cern.ch
 */

#include "config.h"
#include "encryption.h"

#include <string>
#include <iostream>
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <openssl/err.h>

#include <cstdio>
#include <cstring>
#include <cassert>

extern "C" {
   #include "smalloc.h"
}

using namespace std;

namespace encryption {

   unsigned char key[KEY_LENGTH];
   bool key_loaded;
   
   void init() {
      key_loaded = false;
      OpenSSL_add_all_algorithms();
   }
   
   void fini() {
      EVP_cleanup();
      memset(key, 0, KEY_LENGTH);
   }
   
   
   string bin2hex(const unsigned char *buf, const size_t len) {
      string result(2*len, 0);
      for(size_t i = 0; i < len; ++i) {
         char dgt1 = (unsigned)buf[i] / 16;
         char dgt2 = (unsigned)buf[i] % 16;
         dgt1 += (dgt1 <= 9) ? '0' : 'a' - 10;
         dgt2 += (dgt2 <= 9) ? '0' : 'a' - 10;
         result[2*i] = dgt1;
         result[2*i+1] = dgt2;
      }
      return result;
   }
   
   
   void hex2bin(const string hex, unsigned char *buf, const size_t len) {
      assert(hex.length() == 2*len);
      
      for (size_t i = 0; i < 2*len; i+=2) {
         const unsigned char hi = (hex[i] >= 'a') ? hex[i]-'a'+10 : hex[i]-'0';
         const unsigned char lo = (hex[i+1] >= 'a') ? hex[i+1]-'a'+10 : hex[i+1]-'0';
         buf[i/2] = hi*16 + lo;
      }
   }
   

   static bool load_keyfile_internal(FILE *fkey) {
      char buf[2*KEY_LENGTH];
      int bytes_read;
      
      /* 32 bytes key for AES256 */
      bytes_read = fread(buf, 1, 2*KEY_LENGTH, fkey);
      if (bytes_read != 2*KEY_LENGTH)
         return false;
      
      hex2bin(string(buf, 2*KEY_LENGTH), key, KEY_LENGTH);
      
      return true;
   }
   
   
   void random_iv(unsigned char new_iv[IV_LENGTH]) {
      int rv = RAND_pseudo_bytes(new_iv, IV_LENGTH);
      assert(rv >= 0);
   }
   
   
   bool load_keyfile(const string &keyfile) {
      FILE *fkey = fopen(keyfile.c_str(), "r");
      if (!fkey)
         return false;
      if (!load_keyfile_internal(fkey)) {
         fclose(fkey);
         return false;
      }
      fclose(fkey);
      
      key_loaded = true;
      return true;
   }
   
   bool load_keyfile(FILE *fkeyfile) {
      if (!load_keyfile_internal(fkeyfile))
         return false;
      
      key_loaded = true;
      return true;
   }
   
   
   bool load_key(const string &hexkey) {
      if (hexkey.length() != 2*KEY_LENGTH)
         return false;
         
      hex2bin(hexkey, key, KEY_LENGTH);
      key_loaded = true;
      return true;
   }
   
   
   bool encrypt(FILE *fin, FILE *fout, unsigned char iv[IV_LENGTH]) {
      if (!key_loaded)
         return false;
     
      EVP_CIPHER_CTX ctx;
      EVP_CIPHER_CTX_init(&ctx);
      
      if (!EVP_EncryptInit_ex(&ctx, EVP_aes_256_cbc(), NULL, key, iv))
         goto decrypt_fail;
      
      unsigned char in[4096];
      unsigned char out[4096+EVP_MAX_BLOCK_LENGTH];
      int inl, outl;
      while (!feof(fin)) {
         if ((inl = fread(in, 1, 4096, fin)) < 0)
            goto decrypt_fail;
         if (!inl) continue;
         
         if (!EVP_EncryptUpdate(&ctx, out, &outl, in, inl))
            goto decrypt_fail;
         if (fwrite(out, 1, outl, fout) != (unsigned)outl)
            goto decrypt_fail;
      }
      if (!EVP_EncryptFinal(&ctx, out, &outl))
         goto decrypt_fail;
      if (outl && (fwrite(out, 1, outl, fout) != (unsigned)outl))
         goto decrypt_fail;
      
      EVP_CIPHER_CTX_cleanup(&ctx);
      return true;
   
   decrypt_fail:
      EVP_CIPHER_CTX_cleanup(&ctx);
      return false;
   }
   
   
   bool decrypt(FILE *fin, FILE *fout, unsigned char iv[IV_LENGTH]) {
      if (!key_loaded)
         return false;
     
      EVP_CIPHER_CTX ctx;
      EVP_CIPHER_CTX_init(&ctx);
      
      if (!EVP_DecryptInit_ex(&ctx, EVP_aes_256_cbc(), NULL, key, iv))
         goto decrypt_fail;
      
      unsigned char in[4096];
      unsigned char out[4096+EVP_MAX_BLOCK_LENGTH];
      int inl, outl;
      while (!feof(fin)) {
         if ((inl = fread(in, 1, 4096, fin)) < 0)
            goto decrypt_fail;
         if (!inl) continue;
         
         if (!EVP_DecryptUpdate(&ctx, out, &outl, in, inl))
            goto decrypt_fail;
         if (fwrite(out, 1, outl, fout) != (unsigned)outl)
            goto decrypt_fail;
      }
      if (!EVP_DecryptFinal(&ctx, out, &outl))
         goto decrypt_fail;
      if (outl && (fwrite(out, 1, outl, fout) != (unsigned)outl))
         goto decrypt_fail;
      
      EVP_CIPHER_CTX_cleanup(&ctx);
      return true;
   
   decrypt_fail:
      EVP_CIPHER_CTX_cleanup(&ctx);
      return false;
   }
   
   
   bool decrypt_mem(unsigned char iv[IV_LENGTH], const unsigned char *cipher,
                    const unsigned clen, unsigned char **plain, unsigned *plen)
   {
      if (!key_loaded)
         return false;
         
      EVP_CIPHER_CTX ctx;
      EVP_CIPHER_CTX_init(&ctx);
      *plain = (unsigned char *)smalloc(clen+EVP_MAX_BLOCK_LENGTH);
      
      int processed;
      if (!EVP_DecryptInit_ex(&ctx, EVP_aes_256_cbc(), NULL, key, iv))
         goto decrypt_mem_fail;
      if (!EVP_DecryptUpdate(&ctx, *plain, &processed, cipher, clen))
         goto decrypt_mem_fail;
      *plen = processed;
      if (!EVP_DecryptFinal(&ctx, *plain, &processed))
         goto decrypt_mem_fail;
      *plen = *plen + processed;
      
      EVP_CIPHER_CTX_cleanup(&ctx);
      return true;
  
    decrypt_mem_fail:
      EVP_CIPHER_CTX_cleanup(&ctx);
      free(*plain);
      *plain = NULL;
      *plen = 0;
      return false;
   }
}
